import csv
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
from util import Parameter

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((128, 128)),  # 调整图像大小
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 标准化
])
transform_flip=transforms.Compose([
    transforms.RandomHorizontalFlip(p=Parameter.FLIP_SCALE),
    transforms.RandomVerticalFlip(p=Parameter.FLIP_SCALE),
    transform
])

class VegetableDataset(Dataset):
    def __init__(self, meta_file:str, transform=transform_flip):
        self.transform=transform
        self.data=[]
        with open(meta_file,mode='r', newline='', encoding='utf-8') as fp:
            csv_reader=csv.reader(fp)
            for row in csv_reader:
                self.data.append(row)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_path = self.data[idx][0]
        label = int(self.data[idx][1])
        image = Image.open(img_path).convert('RGB')  # 打开图像并转换为 RGB 格式

        if self.transform:
            image = self.transform(image)

        return image, label
    
# if __name__=="__main__":
#     tmp=RiceDataset("train.csv")